{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2024-06-19T14:19:32.862556900Z",
     "start_time": "2024-06-19T14:19:27.384005700Z"
    }
   },
   "outputs": [],
   "source": [
    "import os.path as osp\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "from models.backbones import BackBone\n",
    "from dataloader.samplers import CategoriesSampler\n",
    "from utils import pprint, ensure_path, Averager, count_acc, compute_confidence_interval, \\\n",
    "    calculate_keyPatches_index, calculate_keyPatches_num,emd_inference_opencv_test\n",
    "\n",
    "\n",
    "max_epoch=100\n",
    "way=3\n",
    "test_way=3\n",
    "shot=1\n",
    "query=3\n",
    "step_size=5\n",
    "model_type='small'\n",
    "dataset='miniImageNet'\n",
    "init_weights='./initialization/'\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using small\n",
      "odict_keys(['encoder.cls_token', 'encoder.pos_embed', 'encoder.patch_embed.proj.weight', 'encoder.patch_embed.proj.bias', 'encoder.blocks.0.norm1.weight', 'encoder.blocks.0.norm1.bias', 'encoder.blocks.0.attn.qkv.weight', 'encoder.blocks.0.attn.qkv.bias', 'encoder.blocks.0.attn.proj.weight', 'encoder.blocks.0.attn.proj.bias', 'encoder.blocks.0.norm2.weight', 'encoder.blocks.0.norm2.bias', 'encoder.blocks.0.mlp.fc1.weight', 'encoder.blocks.0.mlp.fc1.bias', 'encoder.blocks.0.mlp.fc2.weight', 'encoder.blocks.0.mlp.fc2.bias', 'encoder.blocks.1.norm1.weight', 'encoder.blocks.1.norm1.bias', 'encoder.blocks.1.attn.qkv.weight', 'encoder.blocks.1.attn.qkv.bias', 'encoder.blocks.1.attn.proj.weight', 'encoder.blocks.1.attn.proj.bias', 'encoder.blocks.1.norm2.weight', 'encoder.blocks.1.norm2.bias', 'encoder.blocks.1.mlp.fc1.weight', 'encoder.blocks.1.mlp.fc1.bias', 'encoder.blocks.1.mlp.fc2.weight', 'encoder.blocks.1.mlp.fc2.bias', 'encoder.blocks.2.norm1.weight', 'encoder.blocks.2.norm1.bias', 'encoder.blocks.2.attn.qkv.weight', 'encoder.blocks.2.attn.qkv.bias', 'encoder.blocks.2.attn.proj.weight', 'encoder.blocks.2.attn.proj.bias', 'encoder.blocks.2.norm2.weight', 'encoder.blocks.2.norm2.bias', 'encoder.blocks.2.mlp.fc1.weight', 'encoder.blocks.2.mlp.fc1.bias', 'encoder.blocks.2.mlp.fc2.weight', 'encoder.blocks.2.mlp.fc2.bias', 'encoder.blocks.3.norm1.weight', 'encoder.blocks.3.norm1.bias', 'encoder.blocks.3.attn.qkv.weight', 'encoder.blocks.3.attn.qkv.bias', 'encoder.blocks.3.attn.proj.weight', 'encoder.blocks.3.attn.proj.bias', 'encoder.blocks.3.norm2.weight', 'encoder.blocks.3.norm2.bias', 'encoder.blocks.3.mlp.fc1.weight', 'encoder.blocks.3.mlp.fc1.bias', 'encoder.blocks.3.mlp.fc2.weight', 'encoder.blocks.3.mlp.fc2.bias', 'encoder.blocks.4.norm1.weight', 'encoder.blocks.4.norm1.bias', 'encoder.blocks.4.attn.qkv.weight', 'encoder.blocks.4.attn.qkv.bias', 'encoder.blocks.4.attn.proj.weight', 'encoder.blocks.4.attn.proj.bias', 'encoder.blocks.4.norm2.weight', 'encoder.blocks.4.norm2.bias', 'encoder.blocks.4.mlp.fc1.weight', 'encoder.blocks.4.mlp.fc1.bias', 'encoder.blocks.4.mlp.fc2.weight', 'encoder.blocks.4.mlp.fc2.bias', 'encoder.blocks.5.norm1.weight', 'encoder.blocks.5.norm1.bias', 'encoder.blocks.5.attn.qkv.weight', 'encoder.blocks.5.attn.qkv.bias', 'encoder.blocks.5.attn.proj.weight', 'encoder.blocks.5.attn.proj.bias', 'encoder.blocks.5.norm2.weight', 'encoder.blocks.5.norm2.bias', 'encoder.blocks.5.mlp.fc1.weight', 'encoder.blocks.5.mlp.fc1.bias', 'encoder.blocks.5.mlp.fc2.weight', 'encoder.blocks.5.mlp.fc2.bias', 'encoder.blocks.6.norm1.weight', 'encoder.blocks.6.norm1.bias', 'encoder.blocks.6.attn.qkv.weight', 'encoder.blocks.6.attn.qkv.bias', 'encoder.blocks.6.attn.proj.weight', 'encoder.blocks.6.attn.proj.bias', 'encoder.blocks.6.norm2.weight', 'encoder.blocks.6.norm2.bias', 'encoder.blocks.6.mlp.fc1.weight', 'encoder.blocks.6.mlp.fc1.bias', 'encoder.blocks.6.mlp.fc2.weight', 'encoder.blocks.6.mlp.fc2.bias', 'encoder.blocks.7.norm1.weight', 'encoder.blocks.7.norm1.bias', 'encoder.blocks.7.attn.qkv.weight', 'encoder.blocks.7.attn.qkv.bias', 'encoder.blocks.7.attn.proj.weight', 'encoder.blocks.7.attn.proj.bias', 'encoder.blocks.7.norm2.weight', 'encoder.blocks.7.norm2.bias', 'encoder.blocks.7.mlp.fc1.weight', 'encoder.blocks.7.mlp.fc1.bias', 'encoder.blocks.7.mlp.fc2.weight', 'encoder.blocks.7.mlp.fc2.bias', 'encoder.blocks.8.norm1.weight', 'encoder.blocks.8.norm1.bias', 'encoder.blocks.8.attn.qkv.weight', 'encoder.blocks.8.attn.qkv.bias', 'encoder.blocks.8.attn.proj.weight', 'encoder.blocks.8.attn.proj.bias', 'encoder.blocks.8.norm2.weight', 'encoder.blocks.8.norm2.bias', 'encoder.blocks.8.mlp.fc1.weight', 'encoder.blocks.8.mlp.fc1.bias', 'encoder.blocks.8.mlp.fc2.weight', 'encoder.blocks.8.mlp.fc2.bias', 'encoder.blocks.9.norm1.weight', 'encoder.blocks.9.norm1.bias', 'encoder.blocks.9.attn.qkv.weight', 'encoder.blocks.9.attn.qkv.bias', 'encoder.blocks.9.attn.proj.weight', 'encoder.blocks.9.attn.proj.bias', 'encoder.blocks.9.norm2.weight', 'encoder.blocks.9.norm2.bias', 'encoder.blocks.9.mlp.fc1.weight', 'encoder.blocks.9.mlp.fc1.bias', 'encoder.blocks.9.mlp.fc2.weight', 'encoder.blocks.9.mlp.fc2.bias', 'encoder.blocks.10.norm1.weight', 'encoder.blocks.10.norm1.bias', 'encoder.blocks.10.attn.qkv.weight', 'encoder.blocks.10.attn.qkv.bias', 'encoder.blocks.10.attn.proj.weight', 'encoder.blocks.10.attn.proj.bias', 'encoder.blocks.10.norm2.weight', 'encoder.blocks.10.norm2.bias', 'encoder.blocks.10.mlp.fc1.weight', 'encoder.blocks.10.mlp.fc1.bias', 'encoder.blocks.10.mlp.fc2.weight', 'encoder.blocks.10.mlp.fc2.bias', 'encoder.blocks.11.norm1.weight', 'encoder.blocks.11.norm1.bias', 'encoder.blocks.11.attn.qkv.weight', 'encoder.blocks.11.attn.qkv.bias', 'encoder.blocks.11.attn.proj.weight', 'encoder.blocks.11.attn.proj.bias', 'encoder.blocks.11.norm2.weight', 'encoder.blocks.11.norm2.bias', 'encoder.blocks.11.mlp.fc1.weight', 'encoder.blocks.11.mlp.fc1.bias', 'encoder.blocks.11.mlp.fc2.weight', 'encoder.blocks.11.mlp.fc2.bias', 'encoder.norm.weight', 'encoder.norm.bias', 'encoder.head.weight', 'encoder.head.bias'])\n",
      "odict_keys(['backbone.cls_token', 'backbone.pos_embed', 'backbone.patch_embed.proj.weight', 'backbone.patch_embed.proj.bias', 'backbone.blocks.0.norm1.weight', 'backbone.blocks.0.norm1.bias', 'backbone.blocks.0.attn.qkv.weight', 'backbone.blocks.0.attn.qkv.bias', 'backbone.blocks.0.attn.proj.weight', 'backbone.blocks.0.attn.proj.bias', 'backbone.blocks.0.norm2.weight', 'backbone.blocks.0.norm2.bias', 'backbone.blocks.0.mlp.fc1.weight', 'backbone.blocks.0.mlp.fc1.bias', 'backbone.blocks.0.mlp.fc2.weight', 'backbone.blocks.0.mlp.fc2.bias', 'backbone.blocks.1.norm1.weight', 'backbone.blocks.1.norm1.bias', 'backbone.blocks.1.attn.qkv.weight', 'backbone.blocks.1.attn.qkv.bias', 'backbone.blocks.1.attn.proj.weight', 'backbone.blocks.1.attn.proj.bias', 'backbone.blocks.1.norm2.weight', 'backbone.blocks.1.norm2.bias', 'backbone.blocks.1.mlp.fc1.weight', 'backbone.blocks.1.mlp.fc1.bias', 'backbone.blocks.1.mlp.fc2.weight', 'backbone.blocks.1.mlp.fc2.bias', 'backbone.blocks.2.norm1.weight', 'backbone.blocks.2.norm1.bias', 'backbone.blocks.2.attn.qkv.weight', 'backbone.blocks.2.attn.qkv.bias', 'backbone.blocks.2.attn.proj.weight', 'backbone.blocks.2.attn.proj.bias', 'backbone.blocks.2.norm2.weight', 'backbone.blocks.2.norm2.bias', 'backbone.blocks.2.mlp.fc1.weight', 'backbone.blocks.2.mlp.fc1.bias', 'backbone.blocks.2.mlp.fc2.weight', 'backbone.blocks.2.mlp.fc2.bias', 'backbone.blocks.3.norm1.weight', 'backbone.blocks.3.norm1.bias', 'backbone.blocks.3.attn.qkv.weight', 'backbone.blocks.3.attn.qkv.bias', 'backbone.blocks.3.attn.proj.weight', 'backbone.blocks.3.attn.proj.bias', 'backbone.blocks.3.norm2.weight', 'backbone.blocks.3.norm2.bias', 'backbone.blocks.3.mlp.fc1.weight', 'backbone.blocks.3.mlp.fc1.bias', 'backbone.blocks.3.mlp.fc2.weight', 'backbone.blocks.3.mlp.fc2.bias', 'backbone.blocks.4.norm1.weight', 'backbone.blocks.4.norm1.bias', 'backbone.blocks.4.attn.qkv.weight', 'backbone.blocks.4.attn.qkv.bias', 'backbone.blocks.4.attn.proj.weight', 'backbone.blocks.4.attn.proj.bias', 'backbone.blocks.4.norm2.weight', 'backbone.blocks.4.norm2.bias', 'backbone.blocks.4.mlp.fc1.weight', 'backbone.blocks.4.mlp.fc1.bias', 'backbone.blocks.4.mlp.fc2.weight', 'backbone.blocks.4.mlp.fc2.bias', 'backbone.blocks.5.norm1.weight', 'backbone.blocks.5.norm1.bias', 'backbone.blocks.5.attn.qkv.weight', 'backbone.blocks.5.attn.qkv.bias', 'backbone.blocks.5.attn.proj.weight', 'backbone.blocks.5.attn.proj.bias', 'backbone.blocks.5.norm2.weight', 'backbone.blocks.5.norm2.bias', 'backbone.blocks.5.mlp.fc1.weight', 'backbone.blocks.5.mlp.fc1.bias', 'backbone.blocks.5.mlp.fc2.weight', 'backbone.blocks.5.mlp.fc2.bias', 'backbone.blocks.6.norm1.weight', 'backbone.blocks.6.norm1.bias', 'backbone.blocks.6.attn.qkv.weight', 'backbone.blocks.6.attn.qkv.bias', 'backbone.blocks.6.attn.proj.weight', 'backbone.blocks.6.attn.proj.bias', 'backbone.blocks.6.norm2.weight', 'backbone.blocks.6.norm2.bias', 'backbone.blocks.6.mlp.fc1.weight', 'backbone.blocks.6.mlp.fc1.bias', 'backbone.blocks.6.mlp.fc2.weight', 'backbone.blocks.6.mlp.fc2.bias', 'backbone.blocks.7.norm1.weight', 'backbone.blocks.7.norm1.bias', 'backbone.blocks.7.attn.qkv.weight', 'backbone.blocks.7.attn.qkv.bias', 'backbone.blocks.7.attn.proj.weight', 'backbone.blocks.7.attn.proj.bias', 'backbone.blocks.7.norm2.weight', 'backbone.blocks.7.norm2.bias', 'backbone.blocks.7.mlp.fc1.weight', 'backbone.blocks.7.mlp.fc1.bias', 'backbone.blocks.7.mlp.fc2.weight', 'backbone.blocks.7.mlp.fc2.bias', 'backbone.blocks.8.norm1.weight', 'backbone.blocks.8.norm1.bias', 'backbone.blocks.8.attn.qkv.weight', 'backbone.blocks.8.attn.qkv.bias', 'backbone.blocks.8.attn.proj.weight', 'backbone.blocks.8.attn.proj.bias', 'backbone.blocks.8.norm2.weight', 'backbone.blocks.8.norm2.bias', 'backbone.blocks.8.mlp.fc1.weight', 'backbone.blocks.8.mlp.fc1.bias', 'backbone.blocks.8.mlp.fc2.weight', 'backbone.blocks.8.mlp.fc2.bias', 'backbone.blocks.9.norm1.weight', 'backbone.blocks.9.norm1.bias', 'backbone.blocks.9.attn.qkv.weight', 'backbone.blocks.9.attn.qkv.bias', 'backbone.blocks.9.attn.proj.weight', 'backbone.blocks.9.attn.proj.bias', 'backbone.blocks.9.norm2.weight', 'backbone.blocks.9.norm2.bias', 'backbone.blocks.9.mlp.fc1.weight', 'backbone.blocks.9.mlp.fc1.bias', 'backbone.blocks.9.mlp.fc2.weight', 'backbone.blocks.9.mlp.fc2.bias', 'backbone.blocks.10.norm1.weight', 'backbone.blocks.10.norm1.bias', 'backbone.blocks.10.attn.qkv.weight', 'backbone.blocks.10.attn.qkv.bias', 'backbone.blocks.10.attn.proj.weight', 'backbone.blocks.10.attn.proj.bias', 'backbone.blocks.10.norm2.weight', 'backbone.blocks.10.norm2.bias', 'backbone.blocks.10.mlp.fc1.weight', 'backbone.blocks.10.mlp.fc1.bias', 'backbone.blocks.10.mlp.fc2.weight', 'backbone.blocks.10.mlp.fc2.bias', 'backbone.blocks.11.norm1.weight', 'backbone.blocks.11.norm1.bias', 'backbone.blocks.11.attn.qkv.weight', 'backbone.blocks.11.attn.qkv.bias', 'backbone.blocks.11.attn.proj.weight', 'backbone.blocks.11.attn.proj.bias', 'backbone.blocks.11.norm2.weight', 'backbone.blocks.11.norm2.bias', 'backbone.blocks.11.mlp.fc1.weight', 'backbone.blocks.11.mlp.fc1.bias', 'backbone.blocks.11.mlp.fc2.weight', 'backbone.blocks.11.mlp.fc2.bias', 'backbone.norm.weight', 'backbone.norm.bias', 'head.mlp.0.weight', 'head.mlp.0.bias', 'head.mlp.2.weight', 'head.mlp.2.bias', 'head.mlp.4.weight', 'head.mlp.4.bias', 'head.last_layer.weight_g', 'head.last_layer.weight_v', 'head.last_layer2.weight_g', 'head.last_layer2.weight_v'])\n",
      "dict_keys(['encoder.cls_token', 'encoder.pos_embed', 'encoder.patch_embed.proj.weight', 'encoder.patch_embed.proj.bias', 'encoder.blocks.0.norm1.weight', 'encoder.blocks.0.norm1.bias', 'encoder.blocks.0.attn.qkv.weight', 'encoder.blocks.0.attn.qkv.bias', 'encoder.blocks.0.attn.proj.weight', 'encoder.blocks.0.attn.proj.bias', 'encoder.blocks.0.norm2.weight', 'encoder.blocks.0.norm2.bias', 'encoder.blocks.0.mlp.fc1.weight', 'encoder.blocks.0.mlp.fc1.bias', 'encoder.blocks.0.mlp.fc2.weight', 'encoder.blocks.0.mlp.fc2.bias', 'encoder.blocks.1.norm1.weight', 'encoder.blocks.1.norm1.bias', 'encoder.blocks.1.attn.qkv.weight', 'encoder.blocks.1.attn.qkv.bias', 'encoder.blocks.1.attn.proj.weight', 'encoder.blocks.1.attn.proj.bias', 'encoder.blocks.1.norm2.weight', 'encoder.blocks.1.norm2.bias', 'encoder.blocks.1.mlp.fc1.weight', 'encoder.blocks.1.mlp.fc1.bias', 'encoder.blocks.1.mlp.fc2.weight', 'encoder.blocks.1.mlp.fc2.bias', 'encoder.blocks.2.norm1.weight', 'encoder.blocks.2.norm1.bias', 'encoder.blocks.2.attn.qkv.weight', 'encoder.blocks.2.attn.qkv.bias', 'encoder.blocks.2.attn.proj.weight', 'encoder.blocks.2.attn.proj.bias', 'encoder.blocks.2.norm2.weight', 'encoder.blocks.2.norm2.bias', 'encoder.blocks.2.mlp.fc1.weight', 'encoder.blocks.2.mlp.fc1.bias', 'encoder.blocks.2.mlp.fc2.weight', 'encoder.blocks.2.mlp.fc2.bias', 'encoder.blocks.3.norm1.weight', 'encoder.blocks.3.norm1.bias', 'encoder.blocks.3.attn.qkv.weight', 'encoder.blocks.3.attn.qkv.bias', 'encoder.blocks.3.attn.proj.weight', 'encoder.blocks.3.attn.proj.bias', 'encoder.blocks.3.norm2.weight', 'encoder.blocks.3.norm2.bias', 'encoder.blocks.3.mlp.fc1.weight', 'encoder.blocks.3.mlp.fc1.bias', 'encoder.blocks.3.mlp.fc2.weight', 'encoder.blocks.3.mlp.fc2.bias', 'encoder.blocks.4.norm1.weight', 'encoder.blocks.4.norm1.bias', 'encoder.blocks.4.attn.qkv.weight', 'encoder.blocks.4.attn.qkv.bias', 'encoder.blocks.4.attn.proj.weight', 'encoder.blocks.4.attn.proj.bias', 'encoder.blocks.4.norm2.weight', 'encoder.blocks.4.norm2.bias', 'encoder.blocks.4.mlp.fc1.weight', 'encoder.blocks.4.mlp.fc1.bias', 'encoder.blocks.4.mlp.fc2.weight', 'encoder.blocks.4.mlp.fc2.bias', 'encoder.blocks.5.norm1.weight', 'encoder.blocks.5.norm1.bias', 'encoder.blocks.5.attn.qkv.weight', 'encoder.blocks.5.attn.qkv.bias', 'encoder.blocks.5.attn.proj.weight', 'encoder.blocks.5.attn.proj.bias', 'encoder.blocks.5.norm2.weight', 'encoder.blocks.5.norm2.bias', 'encoder.blocks.5.mlp.fc1.weight', 'encoder.blocks.5.mlp.fc1.bias', 'encoder.blocks.5.mlp.fc2.weight', 'encoder.blocks.5.mlp.fc2.bias', 'encoder.blocks.6.norm1.weight', 'encoder.blocks.6.norm1.bias', 'encoder.blocks.6.attn.qkv.weight', 'encoder.blocks.6.attn.qkv.bias', 'encoder.blocks.6.attn.proj.weight', 'encoder.blocks.6.attn.proj.bias', 'encoder.blocks.6.norm2.weight', 'encoder.blocks.6.norm2.bias', 'encoder.blocks.6.mlp.fc1.weight', 'encoder.blocks.6.mlp.fc1.bias', 'encoder.blocks.6.mlp.fc2.weight', 'encoder.blocks.6.mlp.fc2.bias', 'encoder.blocks.7.norm1.weight', 'encoder.blocks.7.norm1.bias', 'encoder.blocks.7.attn.qkv.weight', 'encoder.blocks.7.attn.qkv.bias', 'encoder.blocks.7.attn.proj.weight', 'encoder.blocks.7.attn.proj.bias', 'encoder.blocks.7.norm2.weight', 'encoder.blocks.7.norm2.bias', 'encoder.blocks.7.mlp.fc1.weight', 'encoder.blocks.7.mlp.fc1.bias', 'encoder.blocks.7.mlp.fc2.weight', 'encoder.blocks.7.mlp.fc2.bias', 'encoder.blocks.8.norm1.weight', 'encoder.blocks.8.norm1.bias', 'encoder.blocks.8.attn.qkv.weight', 'encoder.blocks.8.attn.qkv.bias', 'encoder.blocks.8.attn.proj.weight', 'encoder.blocks.8.attn.proj.bias', 'encoder.blocks.8.norm2.weight', 'encoder.blocks.8.norm2.bias', 'encoder.blocks.8.mlp.fc1.weight', 'encoder.blocks.8.mlp.fc1.bias', 'encoder.blocks.8.mlp.fc2.weight', 'encoder.blocks.8.mlp.fc2.bias', 'encoder.blocks.9.norm1.weight', 'encoder.blocks.9.norm1.bias', 'encoder.blocks.9.attn.qkv.weight', 'encoder.blocks.9.attn.qkv.bias', 'encoder.blocks.9.attn.proj.weight', 'encoder.blocks.9.attn.proj.bias', 'encoder.blocks.9.norm2.weight', 'encoder.blocks.9.norm2.bias', 'encoder.blocks.9.mlp.fc1.weight', 'encoder.blocks.9.mlp.fc1.bias', 'encoder.blocks.9.mlp.fc2.weight', 'encoder.blocks.9.mlp.fc2.bias', 'encoder.blocks.10.norm1.weight', 'encoder.blocks.10.norm1.bias', 'encoder.blocks.10.attn.qkv.weight', 'encoder.blocks.10.attn.qkv.bias', 'encoder.blocks.10.attn.proj.weight', 'encoder.blocks.10.attn.proj.bias', 'encoder.blocks.10.norm2.weight', 'encoder.blocks.10.norm2.bias', 'encoder.blocks.10.mlp.fc1.weight', 'encoder.blocks.10.mlp.fc1.bias', 'encoder.blocks.10.mlp.fc2.weight', 'encoder.blocks.10.mlp.fc2.bias', 'encoder.blocks.11.norm1.weight', 'encoder.blocks.11.norm1.bias', 'encoder.blocks.11.attn.qkv.weight', 'encoder.blocks.11.attn.qkv.bias', 'encoder.blocks.11.attn.proj.weight', 'encoder.blocks.11.attn.proj.bias', 'encoder.blocks.11.norm2.weight', 'encoder.blocks.11.norm2.bias', 'encoder.blocks.11.mlp.fc1.weight', 'encoder.blocks.11.mlp.fc1.bias', 'encoder.blocks.11.mlp.fc2.weight', 'encoder.blocks.11.mlp.fc2.bias', 'encoder.norm.weight', 'encoder.norm.bias'])\n"
     ]
    }
   ],
   "source": [
    "\n",
    "#  init_weights = osp.join( init_weights,  dataset, 'checkpoint1600.pth')\n",
    "init_weights = \"F:/checkpoint/ViT-S-16/miniImageNet/checkpoint1600.pth\"\n",
    "if  dataset == 'miniImageNet':\n",
    "    from dataloader.mini_imagenet import MiniImageNet as Dataset\n",
    "else:\n",
    "    raise ValueError('Non-supported Dataset.')\n",
    "\n",
    "valset = Dataset('val',args=None)\n",
    "val_sampler = CategoriesSampler(valset.label, 500,  test_way,  shot +  query)\n",
    "val_loader = DataLoader(dataset=valset, batch_sampler=val_sampler, pin_memory=True)\n",
    "\n",
    "model = BackBone(args=None)\n",
    "\n",
    "print('Using {}'.format( model_type))\n",
    "\n",
    "# load pre-trained models (no FC weights)\n",
    "model_dict = model.state_dict()\n",
    "print(model_dict.keys())\n",
    "if init_weights is not None:\n",
    "    pretrained_dict = torch.load(init_weights, map_location='cpu')['teacher']\n",
    "    print(pretrained_dict.keys())\n",
    "    pretrained_dict = {k.replace('backbone', 'encoder'): v for k, v in pretrained_dict.items()}\n",
    "    pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}\n",
    "    print(pretrained_dict.keys())\n",
    "    model_dict.update(pretrained_dict)\n",
    "    model.load_state_dict(model_dict)\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    torch.backends.cudnn.benchmark = True\n",
    "    model = model.cuda()"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-06-19T14:19:34.033110900Z",
     "start_time": "2024-06-19T14:19:32.869024200Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2\n",
      "torch.Size([20, 3, 224, 224])\n",
      "torch.Size([5, 6, 197, 197])\n"
     ]
    }
   ],
   "source": [
    "feat_shot = torch.empty(way,197,384)\n",
    "feat_query = torch.empty(way*query, 197, 384)\n",
    "weight_s = torch.empty(way,196)\n",
    "weight_q = torch.empty(way*query,196)\n",
    "with torch.no_grad():\n",
    "    for i, batch in enumerate(val_loader, 1):\n",
    "        print(i)\n",
    "        print(batch[0].shape)\n",
    "        if torch.cuda.is_available():\n",
    "            data, _ = [_.cuda() for _ in batch]\n",
    "        else:\n",
    "            data = batch[0]\n",
    "        # support集的样本个数 k\n",
    "        k = test_way * shot\n",
    "        data_shot, data_query = data[:k], data[k:]\n",
    "        feat_shot, feat_query = model(data_shot, data_query)\n",
    "        # 5-way 1-shot 设置 feat.shape\n",
    "        # torch.Size([5, 197, 384])\n",
    "        # torch.Size([75, 197, 384])\n",
    "        print(model.attentions[0].shape)\n",
    "        weight_s,_,weight_q,_ = calculate_keyPatches_index(model)\n",
    "        # weight_s = torch.from_numpy(weight_s).unsqueeze(1)\n",
    "        # weight_q = torch.from_numpy(weight_q).view(way,query,-1)\n",
    "        # weight_s = weight_s.permute(1,0,2).repeat(way*query,1,1)\n",
    "        # weight_q = weight_q.index_select(1, torch.LongTensor([0, 0, 1, 1, 2, 2]))\n",
    "        # if i ==2:\n",
    "        break"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-06-19T14:19:39.555142200Z",
     "start_time": "2024-06-19T14:19:34.041743700Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "outputs": [],
   "source": [
    "print(weight_s.shape)\n",
    "print(weight_q.shape)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-06-19T14:19:39.579108300Z",
     "start_time": "2024-06-19T14:19:39.560652200Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "outputs": [],
   "source": [
    "# TODO 特征融合计算原型\n",
    "proto = feat_shot[:, 1:] + 2 * (feat_shot[:, 0:1].repeat(1, feat_shot.shape[1] - 1, 1))\n",
    "feat_query = feat_query[:, 1:] + 2 * (feat_query[:, 0:1].repeat(1, feat_query.shape[1] - 1, 1))\n",
    "# proto = proto.view(proto.shape[0]/way, way, -1, -1)  # shot,way,196,196\n",
    "# proto = proto.mean(dim=0)  # way * 196 * 196"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-06-19T14:21:36.585258800Z",
     "start_time": "2024-06-19T14:19:39.576805Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([15, 5, 196, 196])\n"
     ]
    }
   ],
   "source": [
    "from utils import get_similarity_map\n",
    "similarity_matrix = get_similarity_map(proto,feat_query,way=way)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-06-19T14:21:36.619033100Z",
     "start_time": "2024-06-19T14:21:36.590778400Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "num_proto:5\n",
      "num_query:15\n",
      "tensor([[77.7577, 32.8464,  6.1233, 13.3652, 30.2991],\n",
      "        [31.9372, 32.9909,  7.4696, 26.6694, 14.8906],\n",
      "        [26.6184, 30.3395, 19.9671,  6.1619,  1.5162],\n",
      "        [34.5558, 37.5662,  8.0463, 30.9573, 18.4904],\n",
      "        [52.1163, 14.5565, 19.5579, 15.1650, 62.3176],\n",
      "        [41.2041, 17.0949, 10.8824, 22.4630, 24.0139],\n",
      "        [20.5298, 60.1903, 24.6129, 44.7310, -5.6950],\n",
      "        [28.0204, 19.6473, 20.6683, 22.9454, 22.7538],\n",
      "        [35.3322, 35.3524, 13.5006, 27.8227, 26.9313],\n",
      "        [45.2750, 36.8550,  9.9322, 23.3628, 42.0487],\n",
      "        [90.9703, 23.5193, 19.9708, 10.8191, 41.6308],\n",
      "        [22.4232, 43.4614, 16.5333, 25.3881, 13.1324],\n",
      "        [23.7175, 38.0692, 16.3519, 42.8655,  6.2765],\n",
      "        [16.8005, 27.6722, 25.7269, 47.1753, 13.4138],\n",
      "        [25.3224, 23.7951, 13.4728,  3.3091, 67.3885]], dtype=torch.float64)\n"
     ]
    },
    {
     "data": {
      "text/plain": "0.6000000238418579"
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(similarity_matrix.shape)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-06-19T14:21:54.412389300Z",
     "start_time": "2024-06-19T14:21:36.619033100Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "outputs": [],
   "source": [
    "\n",
    "weight_s = torch.from_numpy(weight_s)\n",
    "weight_q = torch.from_numpy(weight_q)\n",
    "\n"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-06-19T14:18:24.535128600Z",
     "start_time": "2024-06-19T14:18:24.415682600Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "outputs": [
    {
     "data": {
      "text/plain": "torch.Size([2, 196])"
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from utils import get_emd_distance"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-06-19T14:18:24.550521600Z",
     "start_time": "2024-06-19T14:18:24.537125200Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'at' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001B[1;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[1;31mNameError\u001B[0m                                 Traceback (most recent call last)",
      "Cell \u001B[1;32mIn[18], line 1\u001B[0m\n\u001B[1;32m----> 1\u001B[0m \u001B[43mat\u001B[49m\u001B[38;5;241m.\u001B[39mmean(dim\u001B[38;5;241m=\u001B[39m\u001B[38;5;241m1\u001B[39m)\u001B[38;5;241m.\u001B[39mshape\n",
      "\u001B[1;31mNameError\u001B[0m: name 'at' is not defined"
     ]
    }
   ],
   "source": [
    "result = get_emd_distance(similarity_map=similarity_matrix,weight_1=weight_s,weight_2=weight_q)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-06-19T14:18:24.617785300Z",
     "start_time": "2024-06-19T14:18:24.553548400Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "result.shape"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "start_time": "2024-06-19T14:18:24.616782Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "pprint(result)"
   ],
   "metadata": {
    "collapsed": false
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
